import os
import torch
import cv2
import numpy as np
import torch.nn.functional as F
from models.model import  model_Unet_resnet50
from tqdm import tqdm
from data_aug import data_transform

model = model_Unet_resnet50
model.load_state_dict(torch.load('../model_save/Unet_resnet50/71.pth'))
save_path = '../submit/unet_res50_0121'


os.makedirs(save_path,exist_ok=True)
test_data_path = '../suichang_round1_test_partA_210120\suichang_round1_test_partA_210120'


for name in tqdm(os.listdir(test_data_path)):
    image_dir = os.path.join(test_data_path,name)
    ## cv2.imread可以设置参数-1，读取4通道数据，或者选用默认参数只读取三通道
    image = cv2.imread(image_dir,-1)
    image = data_transform(image=image)['image']
    pred = model(image[np.newaxis,...].cuda())
    pred = torch.argmax(F.log_softmax(pred, dim=1), dim=1).cpu().numpy()[0]
    ## 注意将0-9转成1-10
    pred = pred + 1
    pred = np.stack((pred,pred,pred),axis=2)
    pred = pred.astype(np.uint8)
    cv2.imwrite(os.path.join(save_path,name.replace('tif','png')),pred)
